#ifndef PNP_SOLVER_H
#define PNP_SOLVER_H

#include <ceres/ceres.h>
#include <Eigen/Dense>
#include <vector>
#include <iostream>
#include "AnalyticReprojectionError.h"
#include "QuaternionManifold.h"

class PnPSolver {
public:
    PnPSolver(double fx, double fy, double cx, double cy) 
        : fx_(fx), fy_(fy), cx_(cx), cy_(cy) {}
    
    bool Solve(const std::vector<Eigen::Vector3d>& landmarks,
               const std::vector<Eigen::Vector2d>& bearings,
               Eigen::Vector3d& initial_translation,
               Eigen::Quaterniond& initial_rotation) {
        
        if (landmarks.size() != bearings.size() || landmarks.empty()) {
            std::cerr << "Error: landmarks and bearings must have same non-zero size" << std::endl;
            return false;
        }
        
        // 将初始旋转转换为四元数（确保是单位四元数）
        Eigen::Quaterniond q = initial_rotation.normalized();
        
        // 优化变量
        double translation[3] = {initial_translation.x(), initial_translation.y(), initial_translation.z()};
        double quaternion[4] = {q.x(), q.y(), q.z(), q.w()}; // Eigen使用 [x,y,z,w] 格式
        
        // 创建问题
        ceres::Problem problem;
        problem.AddParameterBlock(translation, 3);

        // 添加四元数流形约束
        ceres::EigenQuaternionManifold* quaternion_manifold = new ceres::EigenQuaternionManifold(); // x, y, z, w
        // ceres::QuaternionManifold* quaternion_manifold = new ceres::QuaternionManifold(); // w, x, y, z
        problem.AddParameterBlock(quaternion, 4, quaternion_manifold);

        std::vector<Eigen::Vector3d> opt_landmarks = landmarks;
        // 为每个观测添加残差块（使用解析导数）
        for (size_t i = 0; i < landmarks.size(); ++i) {
            ceres::CostFunction* cost_function = 
                new AnalyticReprojectionError(landmarks[i], bearings[i], fx_, fy_, cx_, cy_);
            
            problem.AddResidualBlock(cost_function, 
                                   nullptr, // 使用默认的损失函数（L2范数）
                                   translation, 
                                   quaternion);

        }
        
        // 设置求解选项
        ceres::Solver::Options options;
        
        options.linear_solver_type = ceres::DENSE_SCHUR;
        options.minimizer_progress_to_stdout = false;
        options.max_num_iterations = 1000;
        options.function_tolerance = 1e-8;
        options.gradient_tolerance = 1e-12;
        options.parameter_tolerance = 1e-10;

        // 暂时禁用梯度检查以测试基本功能
        // options.check_gradients = false;
        // options.gradient_check_relative_precision = 1e-6;

        // 使用更保守的优化策略
        options.trust_region_strategy_type = ceres::DOGLEG;
        // options.dogleg_type = ceres::TRADITIONAL_DOGLEG;

        // 求解
        ceres::Solver::Summary summary;
        ceres::Solve(options, &problem, &summary);
        
        std::cout << summary.BriefReport() << std::endl;
        
        std::cout << "Initial cost: " << summary.initial_cost << std::endl;
        std::cout << "Final cost: " << summary.final_cost << std::endl;
        
        // 输出优化结果
        if (summary.IsSolutionUsable()) {
            // 更新结果
            initial_translation = Eigen::Vector3d(translation[0], translation[1], translation[2]);
            
            // 注意：Eigen四元数构造函数的参数顺序是 (w, x, y, z)
            initial_rotation = Eigen::Quaterniond(quaternion[3], quaternion[0], quaternion[1], quaternion[2]);
            initial_rotation.normalize();
            
            std::cout << "Optimization successful!" << std::endl;
            std::cout << "Final translation: " << initial_translation.transpose() << std::endl;
            std::cout << "Final rotation quaternion: " << initial_rotation.coeffs().transpose() << std::endl;
            
            // 计算旋转矩阵和欧拉角（可选）
            Eigen::Matrix3d R_final = initial_rotation.toRotationMatrix();
            std::cout << "Final rotation matrix:\n" << R_final << std::endl;
            
            return true;
        } else {
            std::cerr << "Optimization failed!" << std::endl;
            return false;
        }
    }

private:
    double fx_, fy_, cx_, cy_;
};

#endif // PNP_SOLVER_H